import torch.nn as nn
from transformers import AutoModel
from configuration import config


class BertClassifier(nn.Module):
    def __init__(self, freeze_bert=True):
        super().__init__()
        self.bert = AutoModel.from_pretrained(config.PRETRAINED_DIR / 'bert-base-chinese')
        self.linear = nn.Linear(self.bert.config.hidden_size, config.NUM_CLASSES)
        if freeze_bert:
            for param in self.bert.parameters():
                param.requires_grad = False

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)  # [batch_size, seq_len, hidden_size]
        cls_output = outputs.last_hidden_state[:, 0, :]  # [batch_size, hidden_size]
        output = self.linear(cls_output)  # [batch_size, num_classes]
        return output
